"""
《邢不行-2019新版|Python股票量化投资课程》
author：邢不行
防断更加微信：Eric20150721

根据选股数据，进行选股
"""
import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
from Functions import *
pd.set_option('expand_frame_repr', False)  # 当列太多时不换行
pd.set_option('display.max_rows', 5000)  # 最多显示数据的行数
class BackTest(object):
    def __init__(self, file_path, index_path, c_rate=1.5 / 10000, t_rate=1 / 1000):
        self.file_path = file_path
        self.index_path = index_path
        self.c_rate = c_rate
        self.t_rate = t_rate
        self.df = pd.read_hdf(self.file_path, 'df')
        self.index_data = import_index_data(index_path)
    def cal_return(self):

        self.df = self.df.iloc[1:][:]
        self.df['代码'] += ' '
        self.df['简称'] += ' '
        self.df['日期'] = self.df['交易日期']
        group = self.df.groupby('日期')
        self.select_stock = pd.DataFrame()
        self.select_stock['买入代码'] = group['代码'].sum()
        self.select_stock['买入简称'] = group['简称'].sum()
        # 计算下周期每天的资金曲线
        self.select_stock['选股下周期每天资金曲线'] = group['下周期每天涨跌幅'].apply(lambda x: np.cumprod(np.array(list(x))+1, axis=1).mean(axis=0))

        self.select_stock['选股下周期每天资金曲线'] = self.select_stock['选股下周期每天资金曲线'] * (1 - self.c_rate)  # 计算有不精准的地方

        # 扣除卖出手续费、印花税。最后一天的资金曲线值，扣除印花税、手续费
        self.select_stock['选股下周期每天资金曲线'] = self.select_stock['选股下周期每天资金曲线'].apply(lambda x: list(x[:-1]) + list([x[-1] * (1 - self.c_rate - self.t_rate)]))

        # 计算下周期整体涨跌幅
        self.select_stock['选股下周期涨跌幅'] = self.select_stock['选股下周期每天资金曲线'].apply(lambda x: x[-1] - 1)

        # 计算下周期每天的涨跌幅
        self.select_stock['选股下周期每天涨跌幅'] = self.select_stock['选股下周期每天资金曲线'].apply(lambda x: list(pd.DataFrame([1] + x).pct_change()[0].iloc[1:]))
        # print(self.select_stock['选股下周期每天涨跌幅'].head())
        # 计算整体资金曲线
        self.select_stock.reset_index(inplace=True)
        self.select_stock['资金曲线'] = (self.select_stock['选股下周期涨跌幅'] + 1).cumprod()
        # ===计算选中股票每天的资金曲线
        # 计算每日资金曲线
        self.equity = pd.merge(left=self.index_data, right=self.select_stock[['日期', '买入代码']], on='日期', how='left', sort=True)
        # 将选股结果和大盘指数合并
        self.equity['持有股票代码'] = self.equity['买入代码'].shift()
        self.equity['持有股票代码'].fillna(method='ffill', inplace=True)
        # print(self.equity[-500:-300])
        self.equity.dropna(subset=['持有股票代码'], inplace=True)
        self.equity = self.equity[(self.equity['日期'] <= self.select_stock.iloc[-1][0] + pd.to_timedelta(7, unit='D')) & (
                    self.equity['日期'] >= self.select_stock.iloc[0][0])]
        del self.equity['买入代码']
        self.equity['涨跌幅'] = self.select_stock['选股下周期每天涨跌幅'].sum()
        self.equity['equity_curve'] = (self.equity['涨跌幅'] + 1).cumprod()
        print('选股最终收益率', self.equity['equity_curve'].iloc[-1] - 1)
        self.equity['benchmark'] = (self.equity['指数涨跌幅'] + 1).cumprod()
        print('上证指数最终收益率', self.equity['benchmark'].iloc[-1] - 1)
        self.equity.set_index('日期', inplace=True)
        del self.select_stock['选股下周期每天资金曲线']

    def cal_sharpe(self):
        #计算夏普比率
        profit_day=self.select_stock['选股下周期每天涨跌幅']
        profit_day=profit_day.reset_index()
        profit_day1=profit_day['选股下周期每天涨跌幅'].sum(axis=0)
        self.sharpe=(np.average(profit_day1)-0.03/365)/np.std(profit_day1)
        print('夏普比率是', self.sharpe)

    def drawplot(self):
        # print(self.equity['指数涨跌幅'])
        # ===画图
        plt.plot(self.equity['equity_curve'])
        plt.plot(self.equity['benchmark'])
        plt.legend(loc='best')
        plt.show()
    def maxdrawdown(self):
        #计算最大回撤
        i = np.argmax((np.maximum.accumulate(self.equity['equity_curve']) - self.equity['equity_curve'])/np.maximum.accumulate(self.equity['equity_curve'])) # end of the period
        j = np.argmax(self.equity['equity_curve'][:i]) # start of period
        self.huiche = (1-self.equity['equity_curve'][i]/self.equity['equity_curve'][j])
        print('最大回撤:', self.huiche)

# file_path='/Users/zhenghuihuang/Desktop/研究生课程/面向金融的python/大作业/xbx_stock_2019_完整代码/data/adjusted_data_all_3.h5'
# index_path='/Users/zhenghuihuang/Library/Mobile Documents/com~apple~CloudDocs/毕业设计/结题报告/毕设设计结题代码+数据/program/选股策略/sz000001.csv'
# BackTest=BackTest(file_path, index_path)
# BackTest.cal_return()
# BackTest.cal_sharpe()
# BackTest.maxdrawdown()
# BackTest.drawplot()
